function [stt,etamat,smooth_st,yfor,vstar,logL,countmat,countobs,countszero,likelvec] = kfilterCumReg(parvec, funcmod, Y, trainvec, solveopt, addsol)
% =========================================================================
% kfilterCumReg 
%
% Model with Mixed Frequency, which makes G and Z deterministically TV, 
% but no Break in the underlying structural matrices
%
% Formerly kfilter_tagtv
% Kalman Filter with Temporal Aggregation, from monthly to quaterly, and G
% time varying
% This is a file for the LMJ project 
% It computes the likelihood a DGSE model where there are time aggregeation
% issues and the model is written in terms of cumulator variables.  
% ========================================================================
%% 
%% DESCRIPTION 
%
%% I. State space representation ot time invariant model 
%
% $$ y_{t} = Z_s s_{t}+ C_sZ $$
% 
% $$ s_{t+1} = T_{t} s_{t} + R_{t} \eta_{t}$$
% 
% $$ V(\eta_{t}) = Q =SDX'SDX $$
%
%% II. Inputs 
% 
% *PAREST*: vector of parameters 
% 
% *FUNCMOD*: model function handle 
% 
% *Y*: [nobs nz] data vector, with NANS 
% 
% *TRAINVEC:* [2 1] vector indicating which row to start in and finish for
% the evaluation of the likelihood. Note: if empty [1 nobs] will be used. 
% 
% *SOLVEOPT*: setting for solution, e.g. for NL solver 
% 
% *ADDSOL*: additional input needed to solve model 
% 
%% III. Construction of time aggregated model 
%
% $$ y_{t} = Z_{t}\alpha_{t} + CZ $$
%
% $$ \alpha_{t} = T_{t}\alpha_{t-1} + R_{t}\eta_{t} $$
% 
% where $$ \alpha_{t}=[s_{t};M_{t}] $$ 
% and $$ M_{t} $$ is a cummulator whose transition equation depends on whether
% aggregating SUMS or AVERAGES 
% 
% $$ M_{t}=\Psi_{t} M_{t} + As_{t} $$
%  
% Since cumulator placed at the bottom of the state, the model solution
% must supply the additional output FIELD, a structure supplied by the solution of 
% the model,that contains structures used to set-up extended state-space 
%
% .NZM      Number of m series 
%
% .NZQ      Number of q series 
%
% .A        [NZQ NS] matrix for the transition of the cumulator 
%            Selects the rows of the state with the elements of the
%           cumulator 
%
% Treatment of the constant changed in July 21 2010. Now the extra rows of
% C for the expanded state space will be created by A*C, hence enter in C
% the constant of those observations that will be part of the cumulator.
% Field CADD no longer needed. 
%
% .FLAG_TAGSUM      == 1 if aggregation is for sums 
%                   == 0 if aggregation is for averages 
% 
%% IV. Construction of time varying matrices 
% Current code defines 3 type of matrices 
% $$ G_{t}, R_{t} $$ that vary deterministacally with 
% month j=1,2,3. 
% 
% The observation matrix $$ Z_{t}=W_{t}*Z $$, where $$ W_{t} $$ is
% extracted every period by removing NANs from the data. 
% Note, however, that the code assumes observations are either observed
% every period or every 3 periods. 
% 
% 
%% V. Output 
% 
% *stt*: [nobs ns] vector of one-sided states 
%
% *etamat*: [nobs nx] vector of smooth disturbances 
%
% *smooth_st*: [nobs ns] vector of smooth states 
% 
% *yfor*: [nobs nz] vector of one step ahead forecasts for observables. 
% Note: notation is not exactly as in data unless NANs all grouped at the 
% end of the observation vector
% 
% *vfor*: [nobs nz] vector of one step ahead forecast errors for observables. 
% Note: same as yfor, notation is not exactly as in data unless NANs all grouped at the 
% end of the observation vector
%
% *logL*: [1 1] log likelihood with constant included 
% 
% *countmat*: [nobs ns nx] decomposition of all states when 1 shock at a
% time turned on. 1 shock per page. 
%
% *counobs*: [nobs nz nx] decomposition of all obserbvables when 1 shock at a
% time turned on. 1 shock per page. 
% 
% *countszero*: [ns nx] decomposition of state S(1) when 1 shock at a time
% turned on. 1 shock per column. 
% 
% *likelvec*: [nobs 1] vector of likelihoods 
%  
%% VI. Additional Notes 
% a) It is assumed the number of periods is divisible by 3 and that the first period 
% is month 1, such that the quarterly series are not observed. 
%
% b) Data are demeaned inside the code  
% 
%% VII. Related codes 
%
% *KFILTER_TAGIN.m* Assumes a deterministic pattern in the 
% variation of Z, from M to Q. Requires additional outputs from the DSGE. 
%
% *KFILTER_TAGIN_FLEX.m* Does not require deterministic pattern in
% variation of Z. Works with any model with NANS and otherwise time
% invariant matrices. 
%
% File Created on 06/25/2010
% Modified        07/01/2010
% 7/21/2010      Changed treatment of constant 
% ========================================================================
%% 
%% ELEMENTS OF THE CODE 

%% 1. Solution of the model
%     FIELD contains structures used to set-up extended state-space 
[G,R,C,eu,SDX,Z,structOne,ssvec,structTwo]=feval(funcmod,parvec,solveopt,addsol);
if ~isequal(eu,[1;1]);return;end

% Note: Above change momentarily for debugging of kfilterCumSplit: RAB
% should remove junk,junk from output if there...


% =========================================================================
%% 2. Set-up extended states space 
%     See buildState.m by AB for an alternative 

%% 2.1 Dimensions and time invariant matrices 
[T, ny]=size(Y); 
ns       =size(G,1); 

if size(Z,1)~=structOne.nzHF;error('Rows in Z must match number of variables observed in High frequency'); end 
nx  = size(SDX,2);

Ztilda  = [Z zeros(structOne.nzHF,structOne.nzLF); ...
    zeros(structOne.nzLF,ns) eye(structOne.nzLF)];
% Selection matrix 
A   = structOne.A;
Ctilda   = [C; structOne.A*C]; 
% Number of periods of High frequency per 1 in Low frequency 

 % type of temporal aggregation
if structOne.flag_tagsum == 1
    sigma =[0 ones(1,structOne.NUMIT-1)];
else
    sigma =(1:structOne.NUMIT);
end

%% 2.2 Page matrices with pages corresponding to quarters {1,2,3,4} 
Tpag  = zeros(ns+structOne.nzLF,ns+structOne.nzLF,structOne.NUMIT);
Rpag  = zeros(ns+structOne.nzLF, nx ,structOne.NUMIT);
for ii=1:structOne.NUMIT 
    if structOne.flag_tagsum == 1
        Tpag(:,:,ii) = [G zeros(ns,structOne.nzLF); ...
            structOne.A*G eye(structOne.nzLF)*sigma(ii)];
        Rpag(:,:,ii) = [R; A*R];
    else
        Tpag(:,:,ii) = [G zeros(ns,structOne.nzLF); ...
            (1/sigma(ii))*structOne.A*G eye(structOne.nzLF)*(sigma(ii)-1)/sigma(ii)];
        Rpag(:,:,ii) = [R; (1/sigma(ii))*structOne.A*R];
    end 
end 

% =========================================================================
%% 3. Demean data      
if any(C~=0)==1;Y=Y-repmat((Ztilda*Ctilda)',[T 1]);end;
Y=Y';

% =========================================================================
%% 4. Initialization
%  Begin with Tpag(:,:,1), Rpag(:,:,1) which should be ZERO for the cummulator 

Pzero=lyapunov_symm(Tpag(:,:,1),Rpag(:,:,1)*(SDX'*SDX)*(Rpag(:,:,1)'));
   
%% 5. Forward filter 
% 
% 5.1 Storage matrices and indicator
vstar=nan(ny,T); 
finmat=nan(ny,ny,T); 
kpartg=nan(ns+structOne.nzLF,ny,T); 
likelvec=zeros(T,1); 
yfor=nan(ny,T);
% Matrices with one additional entry (initialization)
stt=zeros(ns+structOne.nzLF,T+1);
stt(:,1)=zeros(size(Ztilda,2),1);
ptt=zeros(ns+structOne.nzLF,ns+structOne.nzLF,T+1);  
ptt(:,:,1)=Pzero; 

lht      =zeros(T,1);
casevec =kron( ones(T/structOne.NUMIT,1) , (1:structOne.NUMIT)' );
dimz     =zeros(T,1); 

for ii=1:T;
    % 5.2 Remove NANS from ytt
    % Construct $$ Z_{t}=W_{t}Z $$
    ytt=Y(:,ii);
    W   =eye(ny);
    ind =~isnan(ytt);
    rowt =find(~isnan(ytt));
    Wtt = W((ind==1),:);
    % create Z*,y* from
    Ztt = Wtt*Ztilda;
    ytt = ytt(ind);
    
    dimz(ii)=length(ytt);
    
    yfor(rowt,ii)=Ztt*stt(:,ii);
    % 5.3 Kalman routine
    [stt(:,ii+1),ptt(:,:,ii+1),lht(ii),vvv,fin,kgain]=kf_dk(ytt,Ztt,stt(:,ii),ptt(:,:,ii),Tpag(:,:, casevec(ii) ),...
        Rpag(:,:,casevec(ii))*SDX' );
    vstar(rowt,ii) = vvv;
    finmat(rowt,rowt,ii) = fin;
    kpartg(:,rowt,ii) = kgain;
end

logL =-0.5*(sum(dimz))*log(2*pi)+sum(lht(trainvec(1):trainvec(2))); 

% Truncate First Observation which is the initialization 
stt  =stt(:,2:end); 
ptt  =ptt(:,:,2:end); 
yfor=yfor'; 

%% 6. Backward Filter 
% Obtain the Innovation using a disturbance smoother 
% Use same recursion to obtain STATE smoother 

etamat=zeros(nx,T); 
smooth_st=zeros(ns+structOne.nzLF,T); 
rstar=zeros(ns+structOne.nzLF,1);
rmat=zeros(ns+structOne.nzLF,T);
Q=SDX'*SDX;
% First smooth state
[rstar,etamat(:,end)]=smoothdis(rstar,Q,Rpag(:,:,casevec(end))',(Ztt'),finmat(rowt,rowt,end),...
    zeros(ns+structOne.nzLF),vstar(rowt,end));
smooth_st(:,ii)=stt(:,end)+ptt(:,:,end)*rstar;
rmat(:,end)=rstar;
for ii=(T-1):-1:1;
    % 5.2 Remove NANS from ytt
    % Construct $$ Z_{t}=W_{t}Z $$
    ytt=Y(:,ii);
    W   =eye(ny);
    ind =~isnan(ytt);
    rowt =find(~isnan(ytt));
    Wtt = W((ind==1),:);
    % create Z*,y* from    
    Ztt = Wtt*Ztilda;
    %% NOTE on backward filtering 
    %
    % When dealing with time varying G and R matrices need plug 
    % $$ G_{t+1} $$ and $$ R_{t} $$. Note the difference in timing. 
    %
    [rstar,etamat(:,ii)]=smoothdis(rstar,Q,Rpag(:,:,casevec(ii))',...
        (Ztt'),finmat(rowt,rowt,ii),((Tpag(:,:, casevec(ii+1))-Tpag(:,:, casevec(ii+1))*kpartg(:,rowt,ii)*Ztt)'),...
        vstar(rowt,ii));
    smooth_st(:,ii)=stt(:,ii)+ptt(:,:,ii)*rstar;
    rmat(:,ii)=rstar;
end
a0 = Pzero*Tpag(:,:,casevec(1))'*rstar;   % Note: this is only correct in the case where a0 = zeros(ns,1) 
etamat=etamat';
smooth_st=smooth_st';
vstar=vstar';

%% 7. Check Smoother 
% Check that Smooth States are identical if using disturbance smoother (above) vs. state smoother (below) 
% and if can also recover the observables
tol=1e-5; 
% State at time zero give by GG1*azero+Pzero*rstar;
[yCheck,smoothCheck]=kfilterRegSplitSimulation(a0,etamat'); 
disp('Max Discrepancy Smooth vs. Actual Data'); 
maxdifY=comparemat(yCheck,Y'); 
disp('Max Discrepancy State and Innovation Smoother'); 
maxdifS=comparemat(smoothCheck,smooth_st); 
if maxdifY > tol || maxdifS > tol 
    error('Smoother discrepancy exceeds tolerance') 
end 

%% 8. Initial Variance and State per shock (used below for the counterfactual decompositions)
% vInitialPershock: Initial Variance, decomposed per shock 
% sOnePerShock, State at time 1 decomposed, decomposed per shock
vInitialPerShock=zeros(ns,ns,nx);
sZeroSmoothPerShock=zeros(ns,nx); 
sOneSmoothPerShock =zeros(ns,nx); 
for ii=1:nx
    vInitialPerShock(:,:,ii)=lyapunov_symm(Tpag(:,:,casevec(1)),...
        Rpag(:,ii,casevec(1))*SDX(ii,ii)*SDX(ii,ii)*(Rpag(:,ii,casevec(1))'));
    sZeroSmoothPerShock(:,ii) = vInitialPerShock(:,:,ii)*Tpag(:,:,casevec(1))'*rstar;
    etatemp = zeros(nx,1);
    etatemp(ii) = etamat(1,ii);
    sOneSmoothPerShock(:,ii)=Tpag(:,:,casevec(1))*sZeroSmoothPerShock(:,ii)...
        +Rpag(:,:,casevec(1))*etatemp;
end
maxdifSzero=comparemat(sum(sZeroSmoothPerShock,2),a0); 
maxdifSone =comparemat(sum(sOneSmoothPerShock,2),smooth_st(1,:)' ); 
if maxdifSzero > tol; error('Cannot decompose sZeroPerShock'); end 
if maxdifSone > tol; error('Cannot decompose sOnePerShock'); end


%% 9. Counterfactual Decomposition 
% Generate states by feeding each shock at a time, ensure that it recovers
% the original state 
disp('Begin Counterfactual')
countStates=zeros(T,ns,nx);
countObs   =zeros(T,ny,nx);
for ii=1:nx 
    etaTemp=zeros(T,nx); 
    etaTemp(:,ii)=etamat(:,ii); 
    [countObs(:,:,ii),countStates(:,:,ii)]=kfilterRegSplitSimulation(sZeroSmoothPerShock(:,ii),etaTemp');
end 
disp('Maximum Discrepancy Counterfactual States and Smooth States') 
maxdifCount=comparemat(sum(countStates,3),smooth_st); 
if maxdifCount > tol; error('Counterfactual States do not recover smooth states'); end 
   
%% Subroutine kfilterRegSplitSimulation allows to simulate the model  
% Inputs 
% sInitial: [ns 1] Initial State 
% innovMat: [nx T] matrix of innovations 
% By being a sub-routine it has access to all variables defined above 
% Be-careful not to use an index (ii,jj) used above or to repeat variable
% names
    function [ySim,sSim]=kfilterRegSplitSimulation(sInitial,innovMat) 
        if ~isequal(size(innovMat),[nx T]) 
            error('Input innovMat must be [nx T]') 
        end 
        sSim=zeros(ns,T);
        ySim=zeros(size(Y));
        sSim(:,1)=Tpag(:,:,casevec(1))*sInitial + ...
            Rpag(:,:,casevec(1))*innovMat(:,1);
        ySim(:,1)=Ztilda*sSim(:,1);
        for kk=2:T;
            sSim(:,kk)=Tpag(:,:,casevec(kk))*sSim(:,kk-1)+...
                Rpag(:,:,casevec(kk))*innovMat(:,kk);
            ySim(:,kk)=Ztilda*sSim(:,kk);
        end
        ySim=ySim'; 
        sSim=sSim'; 
    end

%% End of file
end